# !/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, getopt, os
from distinguish import distinguish_factory
import load_image
import utils
import cv2
import data5.config
import data5.load_image
import data5.distinguish
import data5.preprocessing
import data5.utils


def run_data5(path, o):
    if os.path.isdir(path):
        if o == '':
            print('请输入 -o 选项')
            usage()
            sys.exit()
        lines = []
        g = data5.distinguish.generator_distinguish(data5.config.model_name, data5.config.char_image_width,
                                                    data5.config.char_image_height)

        next(g)

        listdir = os.listdir(path)
        total = len(listdir)
        j = 0
        for filename in listdir:
            yzm = ''
            if not os.path.isdir(os.path.join(path, filename)):
                continue
            imgs = data5.preprocessing.cut_image(data5.preprocessing.preprocessing(
                utils.load_image_by_filename(os.path.join(path, filename, filename + "." + data5.config.suffix))))
            items = []
            for i in range(9):
                items.append(data5.preprocessing.preprocessing(
                    utils.load_image_by_filename(os.path.join(path, filename, str(i) + "." + data5.config.suffix))))

            for img in imgs:
                img = data5.utils.image_padding(img, data5.config.char_image_width, data5.config.char_image_height)
                xsd = []
                for i in range(9):
                    xsd.append(g.send((img, items[i])))
                yzm = yzm + str(xsd.index(min(xsd)))
            lines.append(filename + "," + yzm + "\n")
            percent = float(j) * 100 / float(total)
            sys.stdout.write("%.2f" % percent)
            sys.stdout.write("%\r")
            sys.stdout.flush()
            j += 1
        print("识别完成，正在排序并写入 mappings 文件中...")
        lines.sort()
        with open(o, 'w') as f:
            f.writelines(lines)
        print('成功生成文件：' + o)
    else:
        print('验证码文件夹不存在！')


def usage():
    print(""" A16验证码识别
    Usage:
      # 识别指定文件夹下的验证码
      python run.py -c 1 -o /test/data-1/mappings.txt /test/data-1
      # 识别指定的验证码
      python run.py -c 1 /test/data-1/0000.jpg
      # 帮助
      python run.py -h|--help 
    Options:
        -c, --category CATEGORY 指定要识别验证码的分类 (1|2|3|4|5)
        -o, --out_mappings_filename MAPPINGS_FILENAME 执行输出的mappings文件名
    """)


def res(yzm, c):
    if c == '1':
        try:
            return yzm + "=" + str(eval(yzm))
        except BaseException:
            return yzm
    elif c == '2' or c == '3':
        return yzm.upper()
    elif c == '4':
        try:
            return str(yzm.index('1'))
        except:
            return '-1'

    elif c == '5':
        pass


def parse_category(c):
    if c == '1':
        cate = 'data1'
    elif c == '2':
        cate = 'data2'
    elif c == '3':
        cate = 'data3'
    elif c == '4':
        cate = 'data4'
    elif c == '5':
        cate = 'data5'
    else:
        print('验证码的分类有误')
        sys.exit()
    return cate


WidthAndHeight = [(80, 350), (60, 200), (80, 200), (45, 150)]


def auto_category(d):
    filename = None
    if os.path.isdir(d):
        l = os.listdir(d)
        for f in l:
            if os.path.splitext(f)[1] == '.jpg':
                filename = os.path.join(d, f)
                break
    elif os.path.isfile(d):
        filename = d
    else:
        print('验证码文件或文件夹不存在！')
        sys.exit()
    if filename is None:
        print(d + ' 文件夹下没有*.jpg的图片')
        sys.exit()
    im = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    s = im.shape
    for i, [w, h] in enumerate(WidthAndHeight):
        if s[0] == w and s[1] == h:
            return str(i + 1)


if __name__ == '__main__':
    opts, args = getopt.getopt(sys.argv[1:], 'hc:o:', ['help', 'category:', 'out_mappings_filename:'])
    category = None
    o = ''
    for opt, value in opts:
        if opt in ('-c', '--category'):
            category = value
        if opt in ('-o', '--out_mappings_filename'):
            o = value
        if opt in ('-h', '--help'):
            usage()

    if len(args) == 0:
        print('请输入待识别验证码的目录(或者单个验证码文件路径)')
        usage()
        sys.exit()

    if category is None:
        category = auto_category(args[0])
        print('您没有使用-c选项指定验证码分类，已经自动判断分类为：' + category)

    cate = parse_category(category)

    if cate == 'data5':
        run_data5(args[0], o)
        sys.exit()
    w, h = WidthAndHeight[int(category) - 1]
    distinguish = distinguish_factory.get(cate)

    if os.path.isdir(args[0]):
        if o == '':
            print('请输入 -o 选项')
            usage()
            sys.exit()
        lines = []
        listdir = os.listdir(args[0])
        gen_full = load_image.generator_full_image_and_label_by_dir(args[0], listdir)
        i = 0
        total = len(listdir)
        for basename, gray_img in gen_full:
            yzm = distinguish.distinguish(gray_img, True)
            yzm = os.path.splitext(basename)[0] + "," + res(yzm, category) + "\n"
            lines.append(yzm)
            percent = float(i) * 100 / float(total)
            sys.stdout.write("%.2f" % percent)
            sys.stdout.write("%\r")
            sys.stdout.flush()
            i += 1
        print("识别完成，正在排序并写入 mappings 文件中...")
        lines.sort()
        with open(o, 'w') as f:
            f.writelines(lines)
        print('成功生成文件：' + o)
    elif os.path.isfile(args[0]):
        yzm = distinguish.distinguish(utils.load_image_by_filename(args[0]), True)
        print("识别的结果为：" + os.path.splitext(os.path.basename(args[0]))[0] + "," + res(yzm, category))
    else:
        print('验证码文件或文件夹不存在！')
